{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "9e24a674",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys,copy,os,inspect\n",
    "if hasattr(sys.modules[__name__], '__file__'):\n",
    "    _file_name = __file__\n",
    "else:\n",
    "    _file_name = inspect.getfile(inspect.currentframe())\n",
    "CURRENT_FILE_PATH = os.path.dirname(_file_name)\n",
    "sys.path.append(os.getcwd()+\"/../neuronVis\")\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import IONData\n",
    "import Visual as nv\n",
    "import BrainRegion\n",
    "from tqdm.notebook import tqdm as tool_bar\n",
    "from distinctipy import distinctipy\n",
    "from io import StringIO\n",
    "from distinctipy import distinctipy\n",
    "from sklearn.cluster import AgglomerativeClustering\n",
    "import seaborn\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "29999fa8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAcwAAAD4CAYAAABhR9aJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAENUlEQVR4nO3YsU3DYBhFUQchGjpEAlSswHD0bMY0FCRBdNBhswG6UhwlRucM8OmXXnEtr6ZpGgCAv12c+gEAsASCCQCBYAJAIJgAEAgmAASCCQCBYAJAIJgAEAgmAASCCQCBYAJAIJgAEAgmAASCCQCBYAJAIJgAEAgmAASXcx26+nl6H4bhbq57HMfbw+u42d/4UDpz38+bcbje22kBHoft+LHa2Or8baeX4f6QA3OOLJYLIJYLIZaLIZaLcXCjDA0AgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEcwZzO+MtjmS3/hxP/QaCr7WdFuJ22tlqGQ5u1GqapjkeAgD/ml+yABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEggkAgWACQCCYABAIJgAEvw8cI+vcHkUCAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#auxilary functions\n",
    "iondata = IONData.IONData()\n",
    "def cluster_visualization(IDs,X,clusters,sampling=0,specific_cluster=-1,colors=[],name='cx'):\n",
    "    clustering = AgglomerativeClustering(clusters).fit(X)\n",
    "    if len(colors)==0:\n",
    "        background_color1 = [0.0,0.0,0.0]\n",
    "        background_color2 = [1.0,1.0,1.0]\n",
    "        colors = distinctipy.get_colors(clusters,[background_color1,background_color2])\n",
    "        distinctipy.color_swatch(colors)\n",
    "    neuronvis = nv.neuronVis()\n",
    "    for i,j in zip(IDs,clustering.labels_):\n",
    "        if specific_cluster!=-1:\n",
    "            if j!=specific_cluster:\n",
    "                continue\n",
    "        if np.random.rand()<sampling:\n",
    "            continue\n",
    "        neuronvis.addNeuronByID(i[0],i[1],somaColor=colors[j],axonColor=colors[j],mirrorToRight=True,dendriteHide=True)\n",
    "    for i in ['vontral','right','anterior']:\n",
    "        neuronvis.render.setView(i)\n",
    "        neuronvis.render.savepng(name+'_'+i+'.png')\n",
    "    neuronvis.render.run()\n",
    "def get_mirrowed_point(p):\n",
    "    y = np.array(p.xyz)\n",
    "    if y[2] > 5700:\n",
    "        y[2] = 11400-y[2]\n",
    "    return y\n",
    "background_color1 = [0.0,0.0,0.0]\n",
    "background_color2 = [1.0,1.0,1.0]\n",
    "colors = distinctipy.get_colors(4,[background_color1,background_color2])\n",
    "distinctipy.color_swatch(colors)\n",
    "br = BrainRegion.BrainRegion()\n",
    "br.praseJson()\n",
    "def get_axonal_length(neuron,regions):\n",
    "    while 1:\n",
    "        try:\n",
    "            neuron_property = iondata.getNeuronPropertyByID(neuron[0],neuron[1])\n",
    "            break\n",
    "        except:\n",
    "            os.remove('../resource/json/'+neuron[0]+'/'+neuron[1]+'.json')\n",
    "    brproperty = BrainRegion.RegionProperty(copy.deepcopy(br))\n",
    "    brproperty.setProperty(neuron_property['projectregion'])\n",
    "    tmp = []\n",
    "    for region in regions:\n",
    "        tmp.append(brproperty.getSumProperty(region))\n",
    "    return tmp\n",
    "projection_regions = {\n",
    "    'ISO':['FRP','SS','MO','VISC','AI','GU','ECT','ACA','AUD','VIS','OLF'],\n",
    "    'CNU':['STRd','STRv','LSX','sAMY','PAL'],\n",
    "    'MB':['SCs','SCm','IC','SNr','VTA','MRN','PAG'],\n",
    "    'P':['PB','PCG','PG','TRN'],\n",
    "    'MY':['SPVC','SPVI','SPVO','GRN','IO','IRN','LRN','MDRN','PARN','PGRN','VNC']\n",
    "              }\n",
    "projection_regions_flatten = []\n",
    "for a in projection_regions:\n",
    "    for b in projection_regions[a]:\n",
    "        projection_regions_flatten.append(b)\n",
    "def cluster_projection_regions_summary(IDs,X,clusters):\n",
    "    tmp = {}\n",
    "    clustering = AgglomerativeClustering(clusters).fit(X)\n",
    "    for i,j in tool_bar(zip(IDs,clustering.labels_)):\n",
    "        if j not in tmp:\n",
    "            tmp[j] = {}\n",
    "        tmp[j][i[0]+i[1]] = get_axonal_length(i,projection_regions_flatten)\n",
    "    tmp2 = []\n",
    "    tmp3 = [0 for i in projection_regions_flatten]\n",
    "    for i in tmp:\n",
    "        for j in tmp[i]:\n",
    "            tmp2.append(tmp[i][j].copy())\n",
    "        tmp2.append(tmp3.copy())\n",
    "    return tmp2,[len(list(tmp[i].keys())) for i in tmp]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "id": "c8d9474c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Scene\n",
    "def save_scene(X,Y,clusters,colors,file_name):\n",
    "    clustering = AgglomerativeClustering(clusters).fit(Y)\n",
    "    tmp = []\n",
    "    for i,j in zip(X,clustering.labels_):\n",
    "        tmp.append({'sampleid':i[0],'name':i[1],'mirror':True,'color':{'r':str(int(colors[j][0]*255)),'g':str(int(colors[j][1]*255)),'b':str(int(colors[j][2]*255))}})\n",
    "    Scene.createScene(tmp,file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fbb4a18",
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron_list = []\n",
    "for i in ['1','2/3','5','6a','6b']:\n",
    "    neuron_list += iondata.getNeuronListBySomaRegion(regionName='PL'+i,fuzzy=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e14b9f46",
   "metadata": {},
   "source": [
    "Neural Networks method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2224be49",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#生成数据\n",
    "if 0:\n",
    "    X = []\n",
    "    Y = []\n",
    "    IDs = []\n",
    "    for neuron in tool_bar(neuron_list):\n",
    "        if neuron['sampleid'][:2] in ['AA','00']:\n",
    "            continue\n",
    "        neuron_tree = iondata.getNeuronTreeByID(neuron['sampleid'],neuron['name'])\n",
    "        IDs.append([neuron['sampleid'],neuron['name']])\n",
    "        for p in neuron_tree.points:\n",
    "            if p.type==3:\n",
    "                continue\n",
    "            if len(p.children)==1:\n",
    "                p_mirrowed = get_mirrowed_point(p)\n",
    "                direction = np.array(get_mirrowed_point(p.children[0]))-p_mirrowed\n",
    "                if np.linalg.norm(direction)==0:\n",
    "                    continue\n",
    "                direction/=np.linalg.norm(direction)\n",
    "                X.append(p_mirrowed.tolist()+[len(IDs)-1])\n",
    "                Y.append(direction)\n",
    "    X = np.array(X)\n",
    "    Y = np.array(Y)\n",
    "    np.save('PL_X',X)\n",
    "    np.save('PL_Y',Y)\n",
    "    np.save('PL_IDs',IDs)\n",
    "else:\n",
    "    X = np.load('PL_X.npy')\n",
    "    Y = np.load('PL_Y.npy')\n",
    "    IDs = np.load('PL_IDs.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c123e6c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#predict problem: [neuron vector + 3d position] ----> next 3d direction\n",
    "import tensorflow as tf\n",
    "class Nueron_Classifier:\n",
    "    def __init__(self,num_dimensions=100,leak=0.1,num_neurons=100):\n",
    "        self.num_dimensions = num_dimensions\n",
    "        self.num_neurons = num_neurons\n",
    "        self.leak = 0.1\n",
    "    def generate(self):\n",
    "        x = tf.keras.layers.Input((4,))\n",
    "        x1_normalized = tf.keras.layers.BatchNormalization()(x[:,:3])\n",
    "        x2_embeddings = tf.keras.layers.Embedding(self.num_neurons,self.num_dimensions,input_length=1)(x[:,3:])\n",
    "        x_merged = tf.concat([x1_normalized,x2_embeddings[:,0,:]],axis=-1)\n",
    "        y = tf.keras.layers.Dense(128,activation='linear')(x_merged)\n",
    "        y = tf.keras.layers.LeakyReLU(alpha=self.leak)(y)\n",
    "        y = tf.keras.layers.BatchNormalization()(y)\n",
    "        y = tf.keras.layers.Dense(256,activation='linear')(y)\n",
    "        y = tf.keras.layers.LeakyReLU(alpha=self.leak)(y)\n",
    "        y = tf.keras.layers.BatchNormalization()(y)\n",
    "        y = tf.keras.layers.Dense(512,activation='linear')(y)\n",
    "        y = tf.keras.layers.LeakyReLU(alpha=self.leak)(y)\n",
    "        y = tf.keras.layers.BatchNormalization()(y)\n",
    "        y = tf.keras.layers.Dense(1024,activation='linear')(y)\n",
    "        y = tf.keras.layers.LeakyReLU(alpha=self.leak)(y)\n",
    "        y = tf.keras.layers.BatchNormalization()(y)\n",
    "        y = tf.keras.layers.Dense(3)(y)\n",
    "        model = tf.keras.Model(x,y)\n",
    "        return model\n",
    "class CustomCallback(tf.keras.callbacks.Callback):\n",
    "    def __init__(self,model,**kwargs):\n",
    "        super(CustomCallback, self).__init__(**kwargs)\n",
    "        self.model = model\n",
    "        self.count = 0\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        self.count+=1\n",
    "        if self.count==5:\n",
    "            self.count = 0\n",
    "            plt.matshow(model.layers[3].weights[0].numpy().T)\n",
    "            plt.show()\n",
    "model = Nueron_Classifier(num_dimensions=150,leak=0.1,num_neurons=len(IDs)).generate()\n",
    "model.compile(loss='mse',optimizer='adam')\n",
    "plt.matshow(model.layers[3].weights[0].numpy().T)\n",
    "epochs = 100\n",
    "history = model.fit(X,Y,epochs=epochs,batch_size=2048,validation_split=0.0,shuffle=True,callbacks=[CustomCallback(model)])\n",
    "np.save('pl_ann_embeddings',model.layers[3].weights[0].numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "711ba501",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#分成4类\n",
    "_X = model.layers[3].weights[0].numpy()\n",
    "plt.matshow(_X.T)\n",
    "plt.show()\n",
    "seaborn.clustermap(_X,col_cluster=False)\n",
    "plt.show()\n",
    "cluster_visualization(IDs,_X,4,0.9,-1,colors,name='PL_ann')\n",
    "for i in range(4):\n",
    "    cluster_visualization(IDs,_X,4,0.9,i,colors,name='PL_ann_'+str(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80ba8ec5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#投射脑区强度图\n",
    "tmp,tmp2 = cluster_projection_regions_summary(IDs,model.layers[3].weights[0].numpy(),4)\n",
    "_Y = np.log(np.array(tmp)+1).T\n",
    "if np.max(_Y)!=0:\n",
    "    _Y /= np.max(_Y)\n",
    "_Y = np.array([_Y*0,_Y,_Y*0]).transpose((1,2,0))\n",
    "for ii,i in enumerate(tmp2):\n",
    "    a = np.sum(tmp2[:ii+1])+ii\n",
    "    _Y[:,a] = _Y[:,a]*0\n",
    "    _Y[:,a,0] = 1\n",
    "plt.matshow(_Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "id": "821969d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#保存场景文件\n",
    "save_scene(IDs,_X,4,colors,'PL_ann.nv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ef9dee0",
   "metadata": {},
   "source": [
    "NBLAST method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d39c8f74",
   "metadata": {},
   "outputs": [],
   "source": [
    "#NBLAST\n",
    "import navis\n",
    "import pandas as pd\n",
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7eab6bb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def readSwc_Navis(neuron):\n",
    "    data=pd.read_csv(StringIO(iondata.getNeuronByID(neuron['sampleid'],neuron['name'])), header=None)\n",
    "    data['node_id'], data['label'], data['x'], data['y'], data['z'], data['radius'],data['parent_id'] = data.iloc[:, 0].str.split(' ').str\n",
    "    data_new=data.drop(labels = 0,axis = 1)\n",
    "    data_new[['x','y','z','radius','label','parent_id','node_id']] = data_new[['x','y','z','radius','label','parent_id','node_id']].apply(pd.to_numeric)\n",
    "    return data_new\n",
    "neuronlist=[]\n",
    "for neuron in tool_bar(neuron_list):\n",
    "    if neuron['sampleid'][:2] in ['AA','00']:\n",
    "        continue\n",
    "    swc_navis=readSwc_Navis(neuron) \n",
    "    n=navis.TreeNeuron(swc_navis)\n",
    "    n.name=neuron['sampleid']+neuron['name']\n",
    "    if n.nodes['z'][0]<5695:\n",
    "        n.nodes[\"z\"]=2*5695-n.nodes[\"z\"]\n",
    "    neuronlist.append(n)\n",
    "lintname=[]\n",
    "for i in range(len(neuronlist)):\n",
    "    lintname.append(neuronlist[i].name)\n",
    "Neurons=reduce(lambda x,y:x+y,neuronlist)\n",
    "nl_um1 = Neurons /100\n",
    "dps=navis.make_dotprops(nl_um1,k=20,resample=False)\n",
    "nb1=navis.nblast_allbyall(dps,progress=False)\n",
    "nb1_mean=(nb1+nb1.T)/2\n",
    "aba_dist = 1-nb1_mean\n",
    "aba_dist.to_csv('PL_nblast.csv')\n",
    "np.save('PL_nblast_IDs',neuronlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad4b046",
   "metadata": {},
   "outputs": [],
   "source": [
    "#分成4类\n",
    "IDs_nblast = []\n",
    "for neuron in tool_bar(neuron_list):\n",
    "    if neuron['sampleid'][:2] in ['AA','00']:\n",
    "        continue\n",
    "    IDs_nblast.append([neuron['sampleid'],neuron['name']])\n",
    "_X = np.array(pd.read_csv('PL_nblast.csv'))[:,1:].astype('float')\n",
    "plt.matshow(_X.T)\n",
    "plt.show()\n",
    "seaborn.clustermap(_X,col_cluster=False)\n",
    "plt.show()\n",
    "cluster_visualization(IDs_nblast,_X,4,0.9,-1,colors,name='PL_nblast')\n",
    "for i in range(4):\n",
    "    cluster_visualization(IDs_nblast,_X,4,0.9,i,colors,name='PL_nblast_'+str(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d5d717",
   "metadata": {},
   "outputs": [],
   "source": [
    "#投射脑区强度图\n",
    "tmp,tmp2 = cluster_projection_regions_summary(IDs_nblast,_X,4)\n",
    "_Y = np.log(np.array(tmp)+1).T\n",
    "if np.max(_Y)!=0:\n",
    "    _Y /= np.max(_Y)\n",
    "_Y = np.array([_Y*0,_Y,_Y*0]).transpose((1,2,0))\n",
    "for ii,i in enumerate(tmp2):\n",
    "    a = np.sum(tmp2[:ii+1])+ii\n",
    "    _Y[:,a] = _Y[:,a]*0\n",
    "    _Y[:,a,0] = 1\n",
    "plt.matshow(_Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e53037",
   "metadata": {},
   "outputs": [],
   "source": [
    "#保存场景文件\n",
    "save_scene(IDs_nblast,_X,4,colors,'PL_nblast.nv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66b73daf",
   "metadata": {},
   "source": [
    "Projection Region method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeff1c81",
   "metadata": {},
   "outputs": [],
   "source": [
    "#分成4类\n",
    "_X = []\n",
    "for i in tool_bar(IDs):\n",
    "    _X.append(get_axonal_length(i,projection_regions_flatten))\n",
    "_X = np.array(_X)\n",
    "_X = np.log(_X+1)\n",
    "if np.max(_X)!=0:\n",
    "    _X /= np.max(_X)\n",
    "plt.matshow(_X.T)\n",
    "plt.show()\n",
    "seaborn.clustermap(_X,col_cluster=False)\n",
    "plt.show()\n",
    "cluster_visualization(IDs,_X,4,0.9,-1,colors,name='PL_pr')\n",
    "for i in range(4):\n",
    "    cluster_visualization(IDs,_X,4,0.9,i,colors,name='PL_pr_'+str(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21e36991",
   "metadata": {},
   "outputs": [],
   "source": [
    "#投射脑区强度图\n",
    "tmp,tmp2 = cluster_projection_regions_summary(IDs,_X,4)\n",
    "_Y = np.log(np.array(tmp)+1).T\n",
    "if np.max(_Y)!=0:\n",
    "    _Y /= np.max(_Y)\n",
    "_Y = np.array([_Y*0,_Y,_Y*0]).transpose((1,2,0))\n",
    "for ii,i in enumerate(tmp2):\n",
    "    a = np.sum(tmp2[:ii+1])+ii\n",
    "    _Y[:,a] = _Y[:,a]*0\n",
    "    _Y[:,a,0] = 1\n",
    "plt.matshow(_Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95806f93",
   "metadata": {},
   "outputs": [],
   "source": [
    "#保存场景文件\n",
    "save_scene(IDs,_X,4,colors,'PL_pr.nv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "438c3eb1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('CX')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "6684bfb8f990107dc2e41da7bb6829210ca36b987584a9d9d8b153337675e487"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
